package com.ailight.smart.ai.listener;

import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;

/**
 * @author: 朱海鹏
 * @Package: com.ailight.smart.ai.listener
 * @Project: ai-light
 * @date: 2025/3/19 17:53
 * @description :
 */

@EqualsAndHashCode(callSuper = true)
@Slf4j
@Data
public class CustomSseListener extends EventSourceListener {

    // 定义错误信息
    private final StringBuilder fullMessage = new StringBuilder();

    private SseEmitter sseEmitter;


    public CustomSseListener(SseEmitter sseEmitter) {
        // 注入SeeMatter
        this.sseEmitter = sseEmitter;
    }

    @Override
    public void onClosed(@NotNull EventSource eventSource) {
        log.info("onClosed: {}", eventSource.request());
        sseEmitter.complete();
    }

    @Override
    public void onEvent(@NotNull EventSource eventSource, @Nullable String id, @Nullable String type, @NotNull String data) {

        try {
            if ("[DONE]".equals(data)) {
                log.info("[DONE]: {}", data);
                // 这个代表发送最终完整消息
                sseEmitter.send(SseEmitter.event()
                        .name("complete")
                        .data(fullMessage.toString()));
                sseEmitter.complete();
                return;
            }
            sseEmitter.send(SseEmitter.event()
            .name("data").data(data));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void onFailure(@NotNull EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
        String errorMsg;
        if (t != null) {
            errorMsg = "API connection failed: " + t.getMessage();
        } else {
            errorMsg = "API connection failed";
        }
        if (response != null) {
            errorMsg += " [HTTP " + response.code() + "]";
        }
        handleError(sseEmitter, errorMsg);
    }

    @Override
    public void onOpen(@NotNull EventSource eventSource, @NotNull Response response) {
        log.info("连接到deep-seek");
    }

    private void handleError(SseEmitter emitter, String message) {
        try {
            emitter.send(SseEmitter.event()
                    .name("error")
                    .data(message));
            emitter.completeWithError(new RuntimeException(message));
        } catch (IOException e) {
            System.err.println("Error sending error event: " + e.getMessage());
        }
    }
}
